from abc import ABC, abstractmethod
import re


class BaseReward(ABC):

    def __init__(self, scope: int):
        self.scope = scope

    @abstractmethod
    def reward(self, generated_texts, reference_texts):
        """
        :param generated_texts: 模型得到的结果
        :param reference_texts: 真正的结果

        return: 得到分数的比例 [-1,1]
        """
        raise NotImplementedError

    @abstractmethod
    def isBreak(self):
        """
        return: 是否阻断流程
        """
        return False


class format_reward(BaseReward):
    def __init__(self, scope: int):
        super().__init__(scope)
        pattern = r'^<｜User｜>(?P<user>.+?)<｜Assistant｜>(?P<assistant>.+?)<｜end｜>.*$'
        self.PATTERN = re.compile(pattern)

    def reward(self, generated_texts, reference_texts):
        # 1. 使用正则匹配,不符合结构为<｜begin_of_sentence｜><｜User｜>xxxxx<｜Assistant｜>XXXXX<｜end｜>,给 -5分并直接返回 其中每个标识符只能存在一次
        generated_match = self.PATTERN.match(generated_texts)
        if not generated_match:
            return -1
        return 0.1

    def isBreak(self):
        return True


class length_reward(BaseReward):
    def __init__(self, scope: int):
        super().__init__(scope)
        pattern = r'^<｜User｜>(?P<user>.+?)<｜Assistant｜>(?P<assistant>.+?)<｜end｜>.*$'
        self.PATTERN = re.compile(pattern)

    def reward(self, generated_texts, reference_texts):
        generated_match = self.PATTERN.match(generated_texts)
        reference_match = self.PATTERN.match(reference_texts)
        # 将match到的拿出来
        generated_assistant = generated_match.group('assistant')
        reference_assistant = reference_match.group('assistant')
        # 2. 数量和原始数量相符不相符 扣分
        if len(generated_assistant) != len(reference_assistant):
            return -1

    def isBreak(self):
        super().isBreak()


class match_reward(BaseReward):
    def __init__(self, scope: int):
        super().__init__(scope)
        pattern = r'^<｜User｜>(?P<user>.+?)<｜Assistant｜>(?P<assistant>.+?)<｜end｜>.*$'
        self.PATTERN = re.compile(pattern)

    def reward(self, generated_texts, reference_texts):
        generated_match = self.PATTERN.match(generated_texts)
        reference_match = self.PATTERN.match(reference_texts)
        # 将match到的拿出来
        generated_assistant = generated_match.group('assistant')
        reference_assistant = reference_match.group('assistant')
        # 3. 对应index相等,不相等,给-2分,相等,给1分
        min_length = min(len(generated_assistant), len(reference_assistant))
        index_reward = 0
        for i in range(min_length):
            if generated_assistant[i] != reference_assistant[i]:
                index_reward -= 1
            else:
                index_reward += 0.5
        index_reward /= min_length
        return index_reward

    def isBreak(self):
        return False



